'''
LangChain_LLM_Math.py, 使用提示词, 转向调用自定义math方法
'''

from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.chains.conversation.memory import ConversationBufferWindowMemory

OPENAI_API_KEY = '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'

llm = ChatOpenAI(openai_api_key=OPENAI_API_KEY, temperature=0, model_name="gpt-3.5-turbo")

from langchain.tools import BaseTool

class EvaluateMathExpression(BaseTool):
    name = "Math Evaluation"
    description = 'use this tool to evaluate a math expression.'

    def _run(self, expr: str):
        return eval(expr)

    def _arun(self, query: str):
        raise NotImplementedError("Async operation not supported yet")


tools = [EvaluateMathExpression()]


agent = initialize_agent(
    agent='chat-conversational-react-description',
    tools=tools,
    llm=llm,
    verbose=True,
    max_iterations=3,
    early_stopping_method='generate',
    memory=ConversationBufferWindowMemory(
        memory_key='chat_history',
        k=5,
        return_messages=True
    )
)


for message in agent.agent.llm_chain.prompt.messages:
    print(message)

agent(f"What is 2 * 2 * 0.13 - 1.001?")


from langchain.agents.conversational_chat.prompt import (PREFIX)
system_message = PREFIX + "\n" + '''
Unfortunately, Assistant is terrible at maths. Assistant should always refers to available tools and never try to answer math questions by itself
'''

new_prompt = agent.agent.create_prompt(
    system_message=system_message,
    tools=tools
)

agent.agent.llm_chain.prompt = new_prompt

agent(f"What is 2 * 2 * 0.13 - 1.001?")